from array import array
from visual import *
import numpy as np

class cubica:
	def __init__(self, t0=0, a=0, b=0, c=0, d=0):
		self.t0 = t0
		self.a = a
		self.b = b
		self.c = c
		self.d = d
	def valor(self, t):
		t = t - self.t0
		return self.a * t**3 + self.b * t**2 + self.c * t + self.d

def tridiagonal(a, b, c, r, n):
	x = array([0] * n, 'd')
		
	for i in range(1, n-1):
		m = a[i] / b[i-1]
		if not m == 0:
			b[i] = b[i] - m * c[i-1]
			r[i] = r[i] - m * r[i-1]
	
	i = n-2
	while 0 <= i < n-1:
		x[i] = (r[i]- c[i] * x[i+1]) / b[i]
		i = i-1
	
	return x

def spline(t, x, n):
	h = array([0] * (n-1), 'd')
	
	for i in range(n-1):
		h[i] = t[i+1] - t[i]
	
	a = array([0] * n, 'd')
	b = array([0] * n, 'd')
	c = array([0] * n, 'd')
	r = array([0] * n, 'd')
	
	a[0] = c[0] = r[0] = 0;
	b[0] = 1;
	a[n-1] = c[n-1] = r[n-1] = 0;
	b[n-1] = 1;
	
	for i in range(1, n-1):
		a[i] = h[i-1]
		b[i] = 2 * (h[i-1] + h[i])
		c[i] = h[i]
		r[i] = 6 * ((x[i+1] - x[i]) / h[i] - (x[i] - x[i-1]) / h[i-1])
	
	s = tridiagonal(a, b, c, r, n)
	
	p = []
		
	for i in range(n-1):
		cub = cubica()
		cub.t0 = t[i]
		cub.a = (s[i+1] - s[i]) / (6 * h[i])
		cub.b = s[i] / 2
		cub.c = (x[i+1] - x[i]) / h[i] - (2 * s[i] + s[i+1]) * h[i] / 6
		cub.d = x[i]
		p.append(cub)
		
	return p

#def plota(t, x, n):	
#	p = spline(t, x, n)
#	c = curve(pos=(0, 0, x[2]), color=color.cyan, radius=.1)
#	
#	for i in range(n-1):
#		for a in arange(t[i], t[i+1], .01):
#			c.append(pos=(a, p[i].valor(a), x[2]))
#

def curva():
	t = array([1., 2., 3.,  4.,  5.,  6.,  7., 8., 9., 10.], 'd')
	x = array([1., 3., 4., -1.,  0., -2., -3., 1., 3.,  2.], 'd')
	y = array([1., 5., 3., -2., -2.,  0.,  1., 2., 0., -1.], 'd')	
	n = 10
	plota(t, x, n)

def plotax(x, y, m, n):
	for j in range(n):
		p = spline(x, m[j], n)
		c = curve(pos=(0, 0, y[j]), color=color.cyan, radius=.1)
		for i in range(n-1):
			for t in arange(x[i], x[i+1], .01):
				c.append(pos=(t, p[i].valor(t), y[j]))

def plotay(x, y, m, n):
	for j in range(n):
		p = spline(x, m[j], n)
		c = curve(pos=(y[j], 0, 0), color=color.cyan, radius=.1)	
		for i in range(n-1):
			for t in arange(x[i], x[i+1], .01):
				c.append(pos=(y[j], p[i].valor(t), t))

def f(x, y):
	return x * sin(x*y)

def superficie():
	n = 10
	x = array(range(n), 'd')
	y = array(range(n), 'd')
	m = np.arange(n*n).reshape(n, n)
	
	for i in range(n):
		for j in range(n):
			m[i][j] = f(x[i], y[j])
	
	plotax(x, y, m, n)
	
	m = m.transpose()
	
#	plotay(y, x, m, n)
			
if __name__ == "__main__":
	i = curve(pos=[(0,0,0), (10,0,0)])
	j = curve(pos=[(0,0,0), (0,10,0)])
	k = curve(pos=[(0,0,0), (0,0,10)])	
	
#	curva()
	superficie()
	
